# -*-coding:utf-8-*-
import glob
import string

"""该文件用于合并kitti数据集中的类，即  将 
    ['Truck','Van','Tram','Car'] 合并为 Car 
    ['Person_sitting','Cyclist','Pedestrian'] 合并为 人
    忽略 Dontcare和Misc类"""


# glob: return document/file path
# 存储Labels文件夹所有txt文件路径，Win下修改路径要注意反斜杠用法，不要用成\
# txt_list = glob.glob('E:/DATASET/vocDataset/training/label_2/*.txt')
txt_list = glob.glob(r'E:\DATASET\VOC_kitti\Labels\*.txt')


# 输出所有文件中的物体类别
def show_category(txt_list):
    category_list = []
    for item in txt_list:
        try:
            with open(item) as tdf:
                for each_line in tdf:
                    labeldata = each_line.strip().split(' ')  # 去掉前后多余的字符并把其分开
                    category_list.append(labeldata[0])  # 只要第一个字段，即类别
        except IOError as ioerr:
            print('File error:' + str(ioerr))
    print(set(category_list))  # 输出集合


# 这里是将每个文件的每行写入line中
def merge(line):
    each_line = ''
    # add space to last word in each line
    for i in range(len(line)):
        if i != (len(line) - 1):
            each_line = each_line + line[i] + ' '
        else:
            each_line = each_line + line[i]  # 最后一条字段后面不加空格
    each_line = each_line + '\n'
    return (each_line)


print('before modify categories are:\n')
# 显示原有的类别
show_category(txt_list)

# item为当前操作的文件，
for item in txt_list:
    new_txt = []
    try:
        # 对item进行读取
        with open(item, 'r') as r_tdf:
            for each_line in r_tdf:
                # 实际上strip是删除的意思；而split则是分割的意思。因此也表示了这两个功能是完全不一样的，
                # strip可以删除字符串的某些字符，而split则是根据规定的字符将字符串进行分割。
                labeldata = each_line.strip().split(' ')

                '''if labeldata[0] in ['Truck','Van','Tram','Car']: # 合并汽车类  
                    labeldata[0] = labeldata[0].replace(labeldata[0],'car')  
                if labeldata[0] in ['Person_sitting','Cyclist','Pedestrian']: # 合并行人类  
                    labeldata[0] = labeldata[0].replace(labeldata[0],'pedestrian')'''
                # print type(labeldata[4])
                # Pedestrian 0.00 0 -0.20 712.40 143.00 810.73 307.92 1.89 0.48 1.20 1.84 1.47 8.41 0.01
                # 没有研究label里数字具体的意思，直接使用前人的代码
                if labeldata[4] == '0.00':
                    labeldata[4] = labeldata[4].replace(labeldata[4], '1.00')
                if labeldata[5] == '0.00':
                    labeldata[5] = labeldata[5].replace(labeldata[5], '1.00')
                # 如果当前行0位置为truck则进行替换，使用小写怕训练时候出错
                if labeldata[0] == 'Truck':
                    labeldata[0] = labeldata[0].replace(labeldata[0], 'truck')
                if labeldata[0] == 'Van':
                    labeldata[0] = labeldata[0].replace(labeldata[0], 'van')
                if labeldata[0] == 'Tram':
                    labeldata[0] = labeldata[0].replace(labeldata[0], 'tram')
                if labeldata[0] == 'Car':
                    labeldata[0] = labeldata[0].replace(labeldata[0], 'car')
                if labeldata[0] in ['Person_sitting', 'Pedestrian']:  # 合并行人类
                    labeldata[0] = labeldata[0].replace(labeldata[0], 'pedestrian')
                if labeldata[0] == 'Cyclist':
                    continue
                if labeldata[0] == 'DontCare':  # 忽略Dontcare类
                    continue
                if labeldata[0] == 'Misc':  # 忽略Misc类
                    continue
                new_txt.append(merge(labeldata))  # 重新写入新的txt文件
        with open(item, 'w+') as w_tdf:  # w+是打开原文件将内容删除，另写新内容进去
            for temp in new_txt:
                w_tdf.write(temp)
    except IOError as ioerr:
        print('File error:' + str(ioerr))

# 显示处理后的种类，这里已经完成了对label的修改
print('after modify categories are: \n')
show_category(txt_list)
